# liner_model.py
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.regression import LinearRegression
from pyspark.ml import Pipeline
from pyspark.sql.functions import lit, col
import os
from pyspark.sql import SparkSession


class SalaryModel:
    def __init__(self, spark_session):
        self.spark = spark_session
        self.pipeline_model = None
        self.is_trained = False

    def train(self, salary_df):
        """训练薪资预测模型"""
        try:
            # 检查是否有足够的数据
            if salary_df.count() < 5:
                raise Exception("数据量不足，无法训练模型")

            # 选择特征列
            feature_cols = ["city", "industry", "education", "experience", "company_size"]

            # 字符串索引器
            indexers = [
                StringIndexer(inputCol=col, outputCol=col + "_index", handleInvalid="keep")
                for col in feature_cols
            ]

            # 特征向量组装
            assembler = VectorAssembler(
                inputCols=[col + "_index" for col in feature_cols],
                outputCol="features"
            )

            # 线性回归模型
            lr = LinearRegression(
                featuresCol="features",
                labelCol="avg_salary",
                maxIter=10,
                regParam=0.3
            )

            # 构建管道
            pipeline = Pipeline(stages=indexers + [assembler, lr])

            # 训练模型
            self.pipeline_model = pipeline.fit(salary_df)
            self.is_trained = True

            print("模型训练成功")
            return self.pipeline_model

        except Exception as e:
            print(f"训练薪资模型失败: {e}")
            raise e

    def save_model(self, model_path):
        """保存训练好的模型"""
        if self.pipeline_model:
            # 确保路径格式正确
            if not model_path.startswith("file://"):
                model_path = "file:///" + os.path.abspath(model_path).replace('\\', '/')

            print(f"保存模型到: {model_path}")
            self.pipeline_model.write().overwrite().save(model_path)
            print(f"模型已保存到: {model_path}")
        else:
            raise Exception("没有训练好的模型可以保存")

    def load_model(self, model_path):
        """加载已训练的模型"""
        try:
            from pyspark.ml import PipelineModel

            # 确保路径格式正确
            if not model_path.startswith("file://"):
                model_path = "file:///" + os.path.abspath(model_path).replace('\\', '/')

            print(f"从路径加载模型: {model_path}")
            self.pipeline_model = PipelineModel.load(model_path)
            self.is_trained = True
            print(f"模型已从 {model_path} 加载")
            return self.pipeline_model
        except Exception as e:
            print(f"加载模型失败: {e}")
            self.is_trained = False
            raise e

    def predict(self, city, industry, education, experience, company_size):
        """使用模型预测薪资"""
        try:
            if not self.pipeline_model or not self.is_trained:
                raise Exception("模型未加载，请先训练或加载模型")

            # 创建测试数据 - 使用lit函数包装字面量
            test_data = [
                (city, industry, education, experience, company_size)
            ]

            schema = ["city", "industry", "education", "experience", "company_size"]
            test_df = self.spark.createDataFrame(test_data, schema)

            # 添加必要的列以保持与训练数据一致
            test_df = test_df.withColumn("position", lit("预测职位")) \
                .withColumn("salary", lit("0-0k")) \
                .withColumn("district", lit("区域")) \
                .withColumn("business_area", lit("商圈")) \
                .withColumn("skills", lit("技能")) \
                .withColumn("company", lit("公司")) \
                .withColumn("min_salary", lit(0.0)) \
                .withColumn("max_salary", lit(0.0)) \
                .withColumn("avg_salary", lit(0.0))

            # 预测
            predictions = self.pipeline_model.transform(test_df)
            prediction_result = predictions.collect()[0]

            # 获取预测值，确保是float类型
            predicted_salary = float(prediction_result["prediction"])

            return max(5.0, round(predicted_salary, 1))  # 确保合理的薪资范围

        except Exception as e:
            print(f"预测错误: {e}")
            import traceback
            traceback.print_exc()
            return 15.0  # 默认值